package org.liuyuantao.lucene.highlight;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.queries.mlt.MoreLikeThis;
import org.apache.lucene.queryparser.classic.MultiFieldQueryParser;
import org.apache.lucene.queryparser.classic.ParseException;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.*;
import org.apache.lucene.search.highlight.*;
import org.apache.lucene.search.vectorhighlight.FastVectorHighlighter;
import org.apache.lucene.search.vectorhighlight.FieldQuery;
import org.apache.lucene.store.Directory;
import org.liuyuantao.lucene.utils.IKAnalyzer5x;

import java.io.IOException;
import java.io.StringReader;

/**
 * 查询结果高亮显示
 * Created by liuyuantao on 2017-7-22.
 */
public class HighlightExample {
    private Directory directory;

    public HighlightExample(Directory directory) {
        this.directory = directory;
    }

    /**
     * 普通查询
     */
    public void normalSearch() {
        try {
            DirectoryReader reader = DirectoryReader.open(this.directory);
            IndexSearcher searcher = new IndexSearcher(reader);
            TermQuery query = new TermQuery(new Term("content", "沈"));
            TopDocs tds = searcher.search(query, 20);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("title"));
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 查询结果高亮
     *
     * @param name
     */
    public void searcherByHighlighter(String name) {
        try {
            Analyzer analyzer = new IKAnalyzer5x(true);
            DirectoryReader reader = DirectoryReader.open(this.directory);
            IndexSearcher searcher = new IndexSearcher(reader);
            MultiFieldQueryParser parser = new MultiFieldQueryParser(new String[]{"name", "content"}, analyzer);
            Query query = parser.parse(name);
            TopDocs tds = searcher.search(query, 20);
            MoreLikeThis mlt = new MoreLikeThis(searcher.getIndexReader());
            mlt.setFieldNames(new String[]{"name", "content"});
            mlt.setMinDocFreq(1);
            mlt.setMinTermFreq(1);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                String title = doc.get("name");
                title = lighterStr(analyzer, query, title, "name");
                System.out.println(title);
                System.out.println("**************************************************************************************************");
                Query moreLike = mlt.like(sd.doc);
                TopDocs stds = searcher.search(moreLike, 10);
                for (ScoreDoc ssd : stds.scoreDocs) {
                    Document d = searcher.doc(ssd.doc);
                    System.out.println(d.get("name"));
                }
            }
            reader.close();
        } catch (IOException | ParseException | InvalidTokenOffsetsException e) {
            e.printStackTrace();
        }
    }

    /**
     * 高亮指定的文本
     *
     * @param analyzer
     * @param query
     * @param txt
     * @param field
     * @return
     * @throws IOException
     * @throws InvalidTokenOffsetsException
     */
    private String lighterStr(Analyzer analyzer, Query query, String txt, String field) throws IOException, InvalidTokenOffsetsException {
        String str = null;
        QueryScorer scorer = new QueryScorer(query);
        Fragmenter fragmenter = new SimpleSpanFragmenter(scorer);
        Formatter fmt = new SimpleHTMLFormatter("<b>", "</b>");
        Highlighter lighter = new Highlighter(fmt, scorer);
        lighter.setTextFragmenter(fragmenter);
        str = lighter.getBestFragments(analyzer.tokenStream(field, new StringReader(txt)), txt, 3, "......\n");
        if (str.equals("")) return txt;
        return str;
    }

    public void searcherByFastHighlighter(String name) {
        try {
            DirectoryReader reader = DirectoryReader.open(this.directory);
            FastVectorHighlighter fvh = new FastVectorHighlighter(false, false);
            Analyzer analyzer = new IKAnalyzer5x(true);
            IndexSearcher searcher = new IndexSearcher(reader);
            MultiFieldQueryParser parser = new MultiFieldQueryParser(new String[]{"name", "content"}, analyzer);
            Query query = parser.parse(name);
            FieldQuery fq = fvh.getFieldQuery(query);
            TopDocs tds = searcher.search(query, 20);
            for (ScoreDoc sd : tds.scoreDocs) {
                String highTitle = fvh.getBestFragment(fq, searcher.getIndexReader(), sd.doc, "name", 100);
                System.out.println(highTitle);
                System.out.println("-------------------------");
                String highContent = fvh.getBestFragment(fq, searcher.getIndexReader(), sd.doc, "content", 100);
                System.out.println(highContent);
                System.out.println("**********************************************************************");
            }
            reader.close();
        } catch (IOException | ParseException e) {
            e.printStackTrace();
        }
    }

    public void lighterDemo() {
        try {
            String txt = "我爱北京天安门，天安门上彩旗飞,伟大领袖毛主席，指引我们向前进，向前进！！！";
            Query query = new QueryParser("f", new IKAnalyzer5x(true)).parse("北京 毛主席");
            QueryScorer scorer = new QueryScorer(query);
            Fragmenter fragment = new SimpleSpanFragmenter(scorer);
            Formatter formatter = new SimpleHTMLFormatter("<span style='color:red'>", "</span>");
            Highlighter highlighter = new Highlighter(formatter, scorer);
            highlighter.setTextFragmenter(fragment);
            String str = highlighter.getBestFragment(new IKAnalyzer5x(true), "f", txt);
            System.out.println(str);
        } catch (IOException | InvalidTokenOffsetsException | ParseException e) {
            e.printStackTrace();
        }
    }
}
